import cv2
import mediapipe as mp
import numpy as np

# 初始化 mediapipe
mp_face_mesh = mp.solutions.face_mesh
face_mesh = mp_face_mesh.FaceMesh(static_image_mode=False, max_num_faces=1)

# 打开摄像头
cap = cv2.VideoCapture(0)

# 读取目标脸（替换用的）
target_img = cv2.imread("face.png")   # 你要换的脸
target_img = cv2.resize(target_img, (480, 640))

# 目标脸关键点（自己量好的）
target_pts = np.array([
    (221, 389),   # 左眼
    (423, 381),  # 右眼
    (325, 593)   # 嘴巴
], dtype=np.float32)

while True:
    ret, frame = cap.read()
    if not ret:
        break

    h, w, _ = frame.shape
    rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    results = face_mesh.process(rgb)

    if results.multi_face_landmarks:
        for face_landmarks in results.multi_face_landmarks:
            # 获取当前帧的关键点
            pts = np.array([
                (int(face_landmarks.landmark[33].x * w), int(face_landmarks.landmark[33].y * h)),  # 左眼
                (int(face_landmarks.landmark[263].x * w), int(face_landmarks.landmark[263].y * h)), # 右眼
                (int(face_landmarks.landmark[13].x * w), int(face_landmarks.landmark[13].y * h))    # 嘴巴
            ], dtype=np.float32)

            # 仿射变换，把目标脸 warp 到当前人脸大小
            M = cv2.getAffineTransform(target_pts, pts)
            warped_face = cv2.warpAffine(target_img, M, (w, h))

            # 创建 mask（全白，表示整张 warped_face 用来融合）
            #mask = 255 * np.ones(warped_face.shape, warped_face.dtype)
            gray = cv2.cvtColor(warped_face, cv2.COLOR_BGR2GRAY)
            _, mask = cv2.threshold(gray, 1, 255, cv2.THRESH_BINARY)

            # 取人脸中心（用鼻尖 landmark）
            nose = face_landmarks.landmark[1]
            cx, cy = int(nose.x * w), int(nose.y * h)

            # -------- 越界检测 + 裁剪 --------
            h_f, w_f = warped_face.shape[:2]
            x1, x2 = cx - w_f // 2, cx + w_f // 2
            y1, y2 = cy - h_f // 2, cy + h_f // 2

            # 裁剪边界
            if x1 < 0:
                warped_face = warped_face[:, -x1:]
                mask = mask[:, -x1:]
                x1 = 0
            if y1 < 0:
                warped_face = warped_face[-y1:, :]
                mask = mask[-y1:, :]
                y1 = 0
            if x2 > w:
                warped_face = warped_face[:, :w - x1]
                mask = mask[:, :w - x1]
                x2 = w
            if y2 > h:
                warped_face = warped_face[:h - y1, :]
                mask = mask[:h - y1, :]
                y2 = h

            # 重新计算融合中心
            new_center = ((x1 + x2) // 2, (y1 + y2) // 2)

            # 融合
            output = cv2.seamlessClone(
                warped_face, frame, mask, new_center, cv2.NORMAL_CLONE
            )

            frame = output

    cv2.imshow("Face Swap", frame)
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()